import json
import os
import time
from selenium import webdriver
from selenium.webdriver.chrome.service import Service


class AutoDriver:
    def __init__(self, url, implicitly_wait=5, cookie=None):
        chrome_options = webdriver.ChromeOptions()

        prefs = {'profile.default_content_settings.popups': 0,  # 防止保存弹窗
                 # 'download.default_directory': DOWN_PATH,  # 设置默认下载路径
                 "profile.default_content_setting_values.automatic_downloads": 1  # 允许多文件下载
                 }
        chrome_options.add_experimental_option('prefs', prefs)

        # 修改windows.navigator.webdriver，防机器人识别机制，selenium自动登陆判别机制
        chrome_options.add_experimental_option('excludeSwitches', ['enable-automation'])
        service = Service("./driver/chromedriver.exe")
        self.url = url,
        self.driver = webdriver.Chrome(service=service, options=chrome_options)
        self.driver.maximize_window()
        self.driver.implicitly_wait(implicitly_wait)
        self.driver.get(url)
        self.cookie = cookie

    def get_cookie(self):
        driver = self.driver
        input("确认登陆成功后Enter键继续...")
        while True:
            with open('.\\driver\\cookies.txt', 'w') as cookie_file:
                cookies = driver.get_cookies()
                for cookie in cookies:
                    if cookie["domain"] not in self.url:  # 校验cookie域名是否正确，否则报错，重新获取
                        print("Cookies is Wrong!")
                        print("cookie校验错误, 30秒后重试")
                        time.sleep(30)
                        driver.refresh()
                        continue
                cookie_file.write(json.dumps(cookies))
                print("cookie 已保存!")
                break

    def set_cookie(self):
        driver = self.driver
        driver.refresh()
        driver.delete_all_cookies()

        if self.cookie is None:
            if not os.path.exists('.\\driver\\cookies.txt'):
                self.get_cookie()
            with open('.\\driver\\cookies.txt', 'r') as cookie:
                # 使用json读取cookies 注意读取的是文件 所以用load而不是loads
                cookies_list = json.load(cookie)
        else:
            cookies_list = json.loads(self.cookie)
        for cookie in cookies_list:
            print(cookie)
            # 该字段有问题所以删除就可以  浏览器打开后记得刷新页面 有的网页注入cookie后仍需要刷新一下
            if 'expiry' in cookie:
                del cookie['expiry']
            if 'sameSite' in cookie:
                cookie['sameSite'] = 'Strict'
            driver.add_cookie(cookie)
        driver.refresh()


if __name__ == '__main__':
    u = "https://www.baidu.com/"
    obj = AutoDriver(u)
